package class06

import "traning/algorithm/utility/heap02"

/*
	给定两个有序数组，arr1，和arr2，再给定一个正数K
	求两个累加和最大的前K个，两个数必须分别来自arr1 和 arr2
	如：	[1,4,5]
        [2,4,6]
	K=2，返回 11（6+5），10（6+4）
	解法：
		最大的数一定时两个数组的最后一个
		把两个数组对应成矩阵，i，j表示arr1的i与arr2j的和，那么下一个最大的数，不是arr[i-1][j] 就是 arr[i][j-1]
			1	4	5
			0	1	2
		2 0	3	6	7
		4 1	5	8	9
		6 2	7	10	11
		初始话大根堆，将最后一个数放入大根堆中
		循环以下流程
			1. 判断大根堆不为空
			2. 弹出大根堆堆顶
			3. 收集答案
			4. 将弹出的节点的左边节点以及上边节点放入大根堆中（使用集合重复节点不放入）
*/
func TopKSum(arr1 []int, arr2 []int, topK int) []int {
	if topK < 0 {
		return []int{}
	}
	// 两个数组的组合最大只有长度的乘积
	topK = min(len(arr1)*len(arr2), topK)
	// 结果集
	ans := make([]int, topK)
	// 填写位置
	index := 0

	// 记录添加过的位置
	set := make([][]bool, len(arr1))
	for i:=0; i<len(set); i++ {
		set[i] = make([]bool, len(arr2))
	}
	// 比较器，初始化为大根堆
	compare := func(node1, node2 interface{}) bool {
		return node1.(MyHeadNode).sum > node2.(MyHeadNode).sum
	}
	mh := heap02.NewHeap(len(arr1)*len(arr2), compare)
	mh.Add(MyHeadNode{len(arr1)-1, len(arr2)-1, arr1[len(arr1)-1]+ arr2[len(arr2)-1]})
	set[len(arr1)-1][len(arr2)-1] = true

	for mh.Size() > 0 && index < topK {
		biggest := mh.Poll().(MyHeadNode)
		ans[index] = biggest.sum
		index++
		if biggest.i - 1 >= 0 && !set[biggest.i - 1][biggest.j] {
			mh.Add(MyHeadNode{biggest.i - 1, biggest.j, arr1[biggest.i - 1] + arr2[biggest.j]})
			set[biggest.i - 1][biggest.j] = true
		}
		if biggest.j - 1 >= 0 && !set[biggest.i][biggest.j-1]{
			mh.Add(MyHeadNode{biggest.i, biggest.j-1, arr1[biggest.i]+ arr2[biggest.j-1]})
			set[biggest.i][biggest.j - 1] = true
		}
	}
	return ans
}


type MyHeadNode struct {
	i   int
	j   int
	sum int
}

